from scipy.spatial.transform import Rotation as R
import glob
import numpy as np
from scipy.optimize import leastsq
import os

def get_Ts_end_to_base(file):
    if not os.path.exists(file):
        print("file not exist!")
        return -1
    with open(file,"r") as f:
        line = f.readline()
        Ts = line.split(" ")
        Ts = [float(i) for i in Ts]
        # 处理欧拉角ABC
        RH = np.array(Ts[3:6])
        R_hand_to_base= R.from_euler('xyz', np.array(RH), degrees=True).as_matrix()
        # 处理平移量XYZ
        T_hand_to_base = Ts[:3]
        # 旋转矩阵和平移量的拼接
        Ts_end_to_base = np.zeros((4, 4), np.float64)
        Ts_end_to_base[:3, :3] = R_hand_to_base
        Ts_end_to_base[:3, 3] = np.array(T_hand_to_base).flatten()
        Ts_end_to_base[3, 3] = 1
        #Ts_end_to_base转置为Ts_base_to_end
        Ts_base_to_end = np.linalg.inv(Ts_end_to_base)
        #圆心坐标点
        Pl = np.ones((4, 1), np.float64)
        point = Ts[6:]
        Pl[:3,0] = point
    return Ts_base_to_end ,Pl


def calibrate(path):
    files_name = glob.glob(os.path.join(path,"*.txt"))
    n_files = len(files_name)
    HB_all = []
    Pl_all = []

    global XS1, XS2, XS3, XS4, XS5, XS6, XS7, XS8
    global T01, T02, T03, T04, T05, T06, T07, T08

    for i in range(n_files):
        file_name = str(i+1)+".txt"
        HB_all.append(get_Ts_end_to_base(os.path.join(path,file_name))[0])
        Pl_all.append(get_Ts_end_to_base(os.path.join(path,file_name))[1])
        # XS1-XS8是线激光坐标系下的球心坐标
        globals()[f"XS{i + 1}"] = (Pl_all[i][:3, 0]).reshape((-1, 1))
        # T01-T08是机器人末端工具坐标系的平移量
        globals()[f"T0{i + 1}"] = (HB_all[i][:3, 3]).reshape((-1, 1))

    XS_difference = np.hstack((XS1 - XS2, XS1 - XS3, XS1 - XS4, XS2 - XS3,
                               XS2 - XS4, XS3 - XS4))
    T0_difference = np.hstack((T02 - T01, T03 - T01, T04 - T01, T03 - T02,
                               T04 - T02, T04 - T03))
    A = XS_difference.T
    # R0是机器人末端工具坐标系的旋转矩阵
    R0 = HB_all[0][:3, :3]
    # R05-R08是机器人末端工具坐标系的旋转矩阵
    R05 = HB_all[4][:3, :3]
    R06 = HB_all[5][:3, :3]
    R07 = HB_all[6][:3, :3]
    R08 = HB_all[7][:3, :3]
    # @表示矩阵乘法
    b = (R0.T @ T0_difference).T
    b1 = (b[:6, 0]).reshape((-1, 1))
    b2 = (b[:6, 1]).reshape((-1, 1))
    b3 = (b[:6, 2]).reshape((-1, 1))

    E = np.array([[0.00000001, 0., 0.],
                  [0., 0.00000001, 0.],
                  [0., 0, 0.00000001]])

    def rotate(X):
        q1 =X[0]
        q2 =X[1]
        q3 =X[2]
        q4 =X[3]
        Rs = np.array([[1 - 2 * (q3 * q3 + q4 * q4), 2 * (q2 * q3 - q1 * q4), 2 * (q1 * q3 + q2 * q4)],
                        [2 * (q2 * q3 + q1 * q4), 1 - 2 * (q2 * q2 + q4 * q4), 2 * (q3 * q4 - q1 * q2)],
                        [2 * (q2 * q4 - q1 * q3), 2 * (q1 * q2 + q3 * q4), 1 - 2 * (q2 * q2 + q3 * q3)]])
        # X是Rt的逆
        X = Rs.T

        X1 = (X[:3, 0]).reshape((-1, 1))
        X2 = (X[:3, 1]).reshape((-1, 1))
        X3 = (X[:3, 2]).reshape((-1, 1))

        X_temp1 = (A.T@A+E)@X1-A.T@b1
        X_temp2 = (A.T@A+E)@X2-A.T@b2
        X_temp3 = (A.T@A+E)@X3-A.T@b3

        return [np.float64(X_temp1[0][0]),np.float64(X_temp1[1][0]),np.float64(X_temp1[2][0])
            ,np.float64(X_temp2[0][0]),np.float64(X_temp2[1][0]),np.float64(X_temp2[2][0])
            ,np.float64(X_temp3[0][0]),np.float64(X_temp3[1][0]),np.float64(X_temp3[2][0])]

    def translation(X):
        t1 =X[0]
        t2 =X[1]
        t3 =X[2]

        Ts = np.array([[t1],[t2],[t3]])

        X_temp1 = ((R06- R05).T@(R06- R05)+E)@Ts - (R06- R05).T@((R05@R@XS5 - R06@R@XS6)+(T05 - T06))
        X_temp2 = ((R07- R06).T@(R07- R06)+E)@Ts - (R07- R06).T@((R06@R@XS6 - R07@R@XS7)+(T06 - T07))
        X_temp3 = ((R08- R07).T@(R08- R07)+E)@Ts - (R08- R07).T@((R07@R@XS7 - R08@R@XS8)+(T07 - T08))
        X_temp4 = ((R07- R05).T@(R07- R05)+E)@Ts - (R07- R05).T@((R05@R@XS5 - R07@R@XS7)+(T05 - T07))
        X_temp5 = ((R08- R05).T@(R08- R05)+E)@Ts - (R08- R05).T@((R05@R@XS5 - R08@R@XS8)+(T05 - T08))
        X_temp6 = ((R08- R06).T@(R08- R06)+E)@Ts - (R08- R06).T@((R06@R@XS6 - R08@R@XS8)+(T06 - T08))

        return [np.float64(X_temp1[0]),np.float64(X_temp1[1]),np.float64(X_temp1[2]),np.float64(X_temp2[0]),np.float64(X_temp2[1]),np.float64(X_temp2[2]),np.float64(X_temp3[0]),np.float64(X_temp3[1]),np.float64(X_temp3[2]),
                np.float64(X_temp4[0]),np.float64(X_temp4[1]),np.float64(X_temp4[2]),np.float64(X_temp5[0]),np.float64(X_temp5[1]),np.float64(X_temp5[2]),np.float64(X_temp6[0]),np.float64(X_temp6[1]),np.float64(X_temp6[2])]
    #求解四元数
    number1 = [0,0,0,0]
    Quaternion = leastsq(rotate, number1)
    R = np.array([[1 - 2 * (Quaternion[0][2] * Quaternion[0][2] + Quaternion[0][3] * Quaternion[0][3]), 2 * (Quaternion[0][1] * Quaternion[0][2] - Quaternion[0][0] * Quaternion[0][3]), 2 * (Quaternion[0][0] * Quaternion[0][2] + Quaternion[0][1] * Quaternion[0][3])],
                    [2 * (Quaternion[0][1] * Quaternion[0][2] + Quaternion[0][0] * Quaternion[0][3]), 1 - 2 * (Quaternion[0][1] * Quaternion[0][1] + Quaternion[0][3] * Quaternion[0][3]), 2 * (Quaternion[0][2] * Quaternion[0][3] - Quaternion[0][0] * Quaternion[0][1])],
                    [2 * (Quaternion[0][1] * Quaternion[0][3] - Quaternion[0][0] * Quaternion[0][2]), 2 * (Quaternion[0][0] * Quaternion[0][1] + Quaternion[0][2] * Quaternion[0][3]), 1 - 2 * (Quaternion[0][1] * Quaternion[0][1] + Quaternion[0][2] * Quaternion[0][2])]])

    # print('四元数:',Quaternion)
    #求解平移向量
    number2 = [0, 0, 0]
    Translation_vector = leastsq(translation, number2)
    RT = np.array([[1 - 2 * (Quaternion[0][2] * Quaternion[0][2] + Quaternion[0][3] * Quaternion[0][3]), 2 * (Quaternion[0][1] * Quaternion[0][2] - Quaternion[0][0] * Quaternion[0][3]), 2 * (Quaternion[0][0] * Quaternion[0][2] + Quaternion[0][1] * Quaternion[0][3]),Translation_vector[0][0]],
                    [2 * (Quaternion[0][1] * Quaternion[0][2] + Quaternion[0][0] * Quaternion[0][3]), 1 - 2 * (Quaternion[0][1] * Quaternion[0][1] + Quaternion[0][3] * Quaternion[0][3]), 2 * (Quaternion[0][2] * Quaternion[0][3] - Quaternion[0][0] * Quaternion[0][1]),Translation_vector[0][1]],
                    [2 * (Quaternion[0][1] * Quaternion[0][3] - Quaternion[0][0] * Quaternion[0][2]), 2 * (Quaternion[0][0] * Quaternion[0][1] + Quaternion[0][2] * Quaternion[0][3]), 1 - 2 * (Quaternion[0][1] * Quaternion[0][1] + Quaternion[0][2] * Quaternion[0][2]),Translation_vector[0][2]],
                   [0,0,0,1]])




    #四元数误差
    WC1 = R @ XS_difference - R0.T @ T0_difference


    #平移向量误差
    Ts = ([Translation_vector[0][0]],[Translation_vector[0][1]],[Translation_vector[0][2]])
    WC2 = []
    WC2.append(((R06 - R05).T @ (R06 - R05) + E) @ Ts - (R06 - R05).T @ ((R05 @ R @ XS5 - R06 @ R @ XS6) + (T05 - T06)))
    WC2.append(((R07 - R06).T @ (R07 - R06) + E) @ Ts - (R07 - R06).T @ ((R06 @ R @ XS6 - R07 @ R @ XS7) + (T06 - T07)))
    WC2.append(((R08 - R07).T @ (R08 - R07) + E) @ Ts - (R08 - R07).T @ ((R07 @ R @ XS7 - R08 @ R @ XS8) + (T07 - T08)))
    WC2.append(((R07 - R05).T @ (R07 - R05) + E) @ Ts - (R07 - R05).T @ ((R05 @ R @ XS5 - R07 @ R @ XS7) + (T05 - T07)))
    WC2.append(((R08 - R05).T @ (R08 - R05) + E) @ Ts - (R08 - R05).T @ ((R05 @ R @ XS5 - R08 @ R @ XS8) + (T05 - T08)))
    WC2.append(((R08 - R06).T @ (R08 - R06) + E) @ Ts - (R08 - R06).T @ ((R06 @ R @ XS6 - R08 @ R @ XS8) + (T06 - T08)))

    ###验证最终圆心
    point_all = []
    for i in range(8):
        point = np.dot(HB_all[i], np.dot(RT, Pl_all[i]))
        point_all.append(point)

    average = [sum(coord) / len(coord) for coord in zip(*point_all)]

    diffs = []

    for coord in point_all:
        diff = [coord[i] - average[i] for i in range(3)]
        diffs.append(diff)

    diffs =np.round(diffs, 3)

    # 打印结果
    print("平均值数组：", average)
    for i, diff in enumerate(diffs):
        print(f"第 {i + 1}次球心坐标误差：", diff)

    res = np.std(point_all, axis=0)  # 计算每一列的标准差
    res2 = np.var(point_all, axis=0)

    TR = np.linalg.inv(RT)

    print('平移向量:',Translation_vector)
    print('四元数残差:',WC1)
    print('Ts残差:',WC2)
    print('圆心数据:',point_all)
    print('圆心标准差:',res)
    print('圆心方差:',res2)
    print('手眼矩阵:', RT)

    return RT, TR, diffs, average

path = r".\data1"
calibrate(path)
